#include <linux/module.h>
#include "gmsmx.h"

void ec_key_print(EC_KEY* key);
void ec_key_gen_test(void);
void sm2_sign_and_verify_test(void);
void external_ec_key_test(void);
void sm2_enc_and_dec_test(void);
void sm2_key_exchang_test(void);

static int gmsmx_init(void)
{
    int ret = 0;
    ret     = mpi_init();
    if (ret) return ret;
    ret = sm2_ec_ctx_init();
    if (ret) return ret;

    ec_key_gen_test();
    sm2_sign_and_verify_test();
    external_ec_key_test();
    sm2_enc_and_dec_test();
    sm2_key_exchang_test();
    log("gmsmx module init\n");
    return 0;
}

static void gmsmx_cleanup(void)
{ 
    sm2_ec_ctx_deinit();
    log("byebye gmsmx module.\n");
}

void ec_key_print(EC_KEY* key)
{
    unsigned char pub[64]  = {0};
    unsigned char priv[32] = {0};
    ec_key_to_buf(key, pub, priv);
    log("EC_KEY:\npubkey:");
    print_hex(pub, 64);
    log("\nprivkey:");
    print_hex(priv, 32);
    log("\n");
}

void ec_key_gen_test()
{
    int     ret;
    EC_KEY* key = ec_key_new();
    log("\n");
    ec_gen_keypair(key);

    // ec_key_print(key);

    ret = ec_check_pubkey(key);
    if (ret) {
        log("pubkey check failed. %d\n", ret);
    } else {
        log("pubkey check ok\n");
    }

    ec_key_free(key);
}

void sm2_sign_and_verify_test()
{
    int           ret = 0;
    const char*   msg = "message digst";
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;
    EC_KEY*       key = ec_key_new();

    log("\n");
    ec_gen_keypair(key);

    ret = ec_check_pubkey(key);
    if (ret) {
        log("pubkey check failed. %d\n", ret);
    } else {
        log("pubkey check ok\n");
    }

    SM2_compute_message_digest((unsigned char*)msg, strlen(msg),
                               (unsigned char*)id, strlen(id), dgst, key);
    // print_hex(dgst, SM3_DIGEST_LENGTH);

    ret = SM2_sign(dgst, &rs, key);
    if (ret) {
        log("sm2 sign failed. %d\n", ret);
    } else {
        log("sm2 sign ok\n");
    }

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_verify(dgst, &rs, key);
    if (ret) {
        log("sm2 verify failed. %d\n", ret);
    } else {
        log("sm2 verify ok\n");
    }
    ec_key_free(key);
}

unsigned char private_1[32] = {0x39, 0x45, 0x20, 0x8f, 0x7b, 0x21, 0x44, 0xb1,
                               0x3f, 0x36, 0xe3, 0x8a, 0xc6, 0xd3, 0x9f, 0x95,
                               0x88, 0x93, 0x93, 0x69, 0x28, 0x60, 0xb5, 0x1a,
                               0x42, 0xfb, 0x81, 0xef, 0x4d, 0xf7, 0xc5, 0xb8};

unsigned char public_1[64] = {
    0x09, 0xF9, 0xDF, 0x31, 0x1E, 0x54, 0x21, 0xA1, 0x50, 0xDD, 0x7D,
    0x16, 0x1E, 0x4B, 0xC5, 0xC6, 0x72, 0x17, 0x9F, 0xAD, 0x18, 0x33,
    0xFC, 0x07, 0x6B, 0xB0, 0x8F, 0xF3, 0x56, 0xF3, 0x50, 0x20, 0xCC,
    0xEA, 0x49, 0x0C, 0xE2, 0x67, 0x75, 0xA5, 0x2D, 0xC6, 0xEA, 0x71,
    0x8C, 0xC1, 0xAA, 0x60, 0x0A, 0xED, 0x05, 0xFB, 0xF3, 0x5E, 0x08,
    0x4A, 0x66, 0x32, 0xF6, 0x07, 0x2D, 0xA9, 0xAD, 0x13,
};

void external_ec_key_test()
{
    int           ret;
    EC_KEY        key;
    const char*   msg = "message digst";
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;

    key.pub    = mpi_point_new(0);
    key.priv   = mpi_read_raw_data(private_1, 32);
    key.pub->x = mpi_read_raw_data(public_1, 32);
    key.pub->y = mpi_read_raw_data(public_1 + 32, 32);
    key.pub->z = mpi_new(0);
    mpi_set_ui(key.pub->z, 1);
    log("\n");

    // ec_key_print(&key);
    ret = ec_check_pubkey(&key);
    if (ret) {
        log("pubkey check failed. %d\n", ret);
    } else {
        log("pubkey check ok\n");
    }

    SM2_compute_message_digest((unsigned char*)msg, strlen(msg),
                               (unsigned char*)id, strlen(id), dgst, &key);
    // print_hex(dgst, SM3_DIGEST_LENGTH);

    ret = SM2_sign(dgst, &rs, &key);
    if (ret) {
        log("sm2 sign failed. %d\n", ret);
    } else {
        log("sm2 sign ok\n");
    }

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_verify(dgst, &rs, &key);
    if (ret) {
        log("sm2 verify failed. %d\n", ret);
    } else {
        log("sm2 verify ok\n");
    }
}

void sm2_enc_and_dec_test()
{
    const char*   in = "encryption standard";
    unsigned char cipher[128];
    unsigned char out[128] = {0};
    size_t        outlen;
    int           ret;

    EC_KEY* key = ec_key_new();
    log("\n");

    ec_gen_keypair(key);

    log("encrypt int:%s, len=%ld\n", in, strlen(in));

    ret =
        SM2_encrypt((const unsigned char*)in, strlen(in), cipher, &outlen, key);
    if (ret) {
        log("sm2 encrypt failed. %d\n", ret);
        goto free;
    } else {
        log("sm2 encrypt ok\n");
    }
    log("cipher(%ld):", outlen);
    print_hex(cipher, outlen);

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_decrypt(cipher, outlen, out, &outlen, key);
    if (ret) {
        log("sm2 decrypt failed. %d\n", ret);
        goto free;
    } else {
        log("sm2 decrypt ok\n");
    }
    log("decrypt out:%s, len=%ld\n", out, outlen);
free:
    ec_key_free(key);
}

#define KEYLEN 16
void sm2_key_exchang_test()
{
    int         rc  = 0;
    const char* Aid = "123456";
    const char* Bid = "123456";

    unsigned char Ra[SM2_PUBKEY_LEN];
    unsigned char Rb[SM2_PUBKEY_LEN];
    unsigned char kab[KEYLEN];
    unsigned char kba[KEYLEN];
    unsigned char suma[SM3_DIGEST_LENGTH];
    unsigned char sumb[SM3_DIGEST_LENGTH];

    SM2_KAP_CTX ctxA;
    SM2_KAP_CTX ctxB;

    EC_KEY* keyA = ec_key_new();
    EC_KEY* keyB = ec_key_new();

    ec_gen_keypair(keyA);
    ec_gen_keypair(keyB);

    rc = SM2_KAP_CTX_init(&ctxA, keyA, keyB->pub, Aid, strlen(Aid), Bid,
                          strlen(Bid), 1, 1);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_CTX_init(&ctxB, keyB, keyA->pub, Bid, strlen(Bid), Aid,
                          strlen(Aid), 0, 1);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_prepare(&ctxA, Ra);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_prepare(&ctxB, Rb);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_compute_key(&ctxA, Rb, kab, KEYLEN, suma);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_compute_key(&ctxB, Ra, kba, KEYLEN, sumb);
    if (rc) {
        log("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_final_check(&ctxA, sumb);
    if (rc) {
        log("error rc = %d\n", rc);
    }
    rc = SM2_KAP_final_check(&ctxB, suma);
    if (rc) {
        log("error rc = %d\n", rc);
    }

    log("key:\n");
    print_hex(kab, KEYLEN);
    print_hex(kba, KEYLEN);
end:
    SM2_KAP_CTX_cleanup(&ctxA);
    SM2_KAP_CTX_cleanup(&ctxB);
    ec_key_free(keyA);
    ec_key_free(keyB);
}

module_init(gmsmx_init);
module_exit(gmsmx_cleanup);
MODULE_LICENSE("GPL v2");